import mlc
import os
import pathlib
current_file_path = pathlib.Path(__file__).parent.resolve()
docker_os = {
    "ubuntu": ["18.04", "20.04", "22.04"],
    "rhel": ["9"]
}
dataset = {
    "resnet50": "imagenet",
    "retinanet": "openimages",
    "bert-99.9": "squad"
}
variations = {
    "resnet50": {
        "tensorflow": {
            "cpu": ["python"]
        },
        "onnxruntime": {
            "cpu": ["python", "cpp"]
        },
        "pytorch": {
            "cpu": []
        }
    },
    "retinanet": {
        "tensorflow": {
        },
        "onnxruntime": {
            "cpu": ["python", "cpp"]
        },
        "pytorch": {
            "cpu": ["python"]
        }
    },
    "bert-99.9": {
        "tensorflow": {
            "cpu": ["python"]
        },
        "onnxruntime": {
            "cpu": ["python"]
        },
        "pytorch": {
            "cpu": []
        }
    }
}

for _os in docker_os:
    for version in docker_os[_os]:
        for model in variations:
            for backend in variations[model]:
                for device in variations[model][backend]:
                    for implementation in variations[model][backend][device]:
                        variation_string = ",_" + model + ",_" + \
                            backend + ",_" + device + ",_" + implementation
                        file_name_ext = "_" + implementation + "_" + backend + "_" + device
                        dockerfile_path = os.path.join(
                            current_file_path,
                            'dockerfiles',
                            model,
                            _os +
                            '_' +
                            version +
                            file_name_ext +
                            '.Dockerfile')
                        mlc_input = {'action': 'run',
                                     'automation': 'script',
                                     'tags': 'app,mlperf,inference,generic' + variation_string,
                                     'adr': {'compiler':
                                             {'tags': 'gcc'},
                                             'inference-src':
                                             {'tags': '_octoml'},
                                             'openimages-preprocessed':
                                             {'tags': '_50'}
                                             },
                                     'print_deps': True,
                                     'quiet': True,
                                     'silent': True,
                                     'fake_run': True
                                     }
                        r = mlc.access(mlc_input)
                        print_deps = r['new_state']['print_deps']
                        comments = ["#RUN " + dep for dep in print_deps]
                        comments.append("")
                        comments.append(
                            "# Run CM workflow for MLPerf inference")
                        mlc_docker_input = {'action': 'run',
                                            'automation': 'script',
                                            'tags': 'build,dockerfile',
                                            'docker_os': _os,
                                            'docker_os_version': version,
                                            'file_path': dockerfile_path,
                                            'comments': comments,
                                            'run_cmd': 'mlc run script --tags=app,mlperf,inference,generic' + variation_string + ' --adr.compiler.tags=gcc --adr.inference-src.tags=_octoml',
                                            'script_tags': 'app,mlperf,inference,generic',
                                            'quiet': True,
                                            'print_deps': True,
                                            'real_run': True
                                            }
                        r = mlc.access(mlc_docker_input)
                        if r['return'] > 0:
                            print(r)
                            exit(1)

                        print('')
                        print("Dockerfile generated at " + dockerfile_path)
